﻿namespace PugLib.Functional
{
	using System;
	using System.Collections;
	using System.Collections.Generic;

	public sealed class MapReduceMaster<TMapKey, TMapValue, TReduceKey, TReduceValue>
	{
		#region Delegates

		public delegate void MapFunction(TMapKey key, TMapValue value, MapReduceList<TReduceKey, TReduceValue> reduceResult);

		public delegate void ReduceFunction(TReduceKey key, IList<TReduceValue> values, IList<TReduceValue> result);

		#endregion

		private readonly MapFunction _mapFunction;
		private readonly ReduceFunction _reduceFunction;


		public MapReduceMaster(MapFunction mapFunction, ReduceFunction reduceFunction)
		{
			if (mapFunction == null)
			{
				throw new ArgumentNullException("mapFunction");
			}
			if (reduceFunction == null)
			{
				throw new ArgumentNullException("reduceFunction");
			}

			_mapFunction = mapFunction;
			_reduceFunction = reduceFunction;
		}

		public IDictionary<TReduceKey, IList<TReduceValue>> Execute(IDictionary<TMapKey, TMapValue> input)
		{
			Dictionary<TReduceKey, List<TReduceValue>> intermediateResult = new Dictionary<TReduceKey, List<TReduceValue>>();

			foreach (KeyValuePair<TMapKey, TMapValue> inputPair in input)
			{
				MapReduceList<TReduceKey, TReduceValue> mapResult = new MapReduceList<TReduceKey, TReduceValue>();

				_mapFunction(inputPair.Key, inputPair.Value, mapResult);

				ProcessMapResult(intermediateResult, mapResult);
			}

			Dictionary<TReduceKey, IList<TReduceValue>> result = new Dictionary<TReduceKey, IList<TReduceValue>>();
			foreach (KeyValuePair<TReduceKey, List<TReduceValue>> intermediatePair in intermediateResult)
			{
				List<TReduceValue> reduceResult = new List<TReduceValue>();

				_reduceFunction(intermediatePair.Key, intermediatePair.Value, reduceResult);

				ProcessReduceResult(result, intermediatePair.Key, reduceResult);
			}
			return result;
		}

		private static void ProcessReduceResult(IDictionary<TReduceKey, IList<TReduceValue>> result,
		                                        TReduceKey key,
		                                        IEnumerable<TReduceValue> reduceResult)
		{
			foreach (TReduceValue value in reduceResult)
			{
				if (!result.ContainsKey(key))
				{
					result[key] = new List<TReduceValue>();
				}
				result[key].Add(value);
			}
		}

		private static void ProcessMapResult(IDictionary<TReduceKey, List<TReduceValue>> intermediate,
		                                     IEnumerable<KeyValuePair<TReduceKey, TReduceValue>> mapResult)
		{
			foreach (KeyValuePair<TReduceKey, TReduceValue> pair in mapResult)
			{
				if (!intermediate.ContainsKey(pair.Key))
				{
					intermediate[pair.Key] = new List<TReduceValue>();
				}
				intermediate[pair.Key].Add(pair.Value);
			}
		}
	}

	public sealed class MapReduceList<TKey, TValue> : IEnumerable<KeyValuePair<TKey, TValue>>
	{
		private readonly List<KeyValuePair<TKey, TValue>> _list = new List<KeyValuePair<TKey, TValue>>();

		#region IEnumerable<KeyValuePair<TKey,TValue>> Members

		public IEnumerator<KeyValuePair<TKey, TValue>> GetEnumerator()
		{
			return _list.GetEnumerator();
		}

		IEnumerator IEnumerable.GetEnumerator()
		{
			return GetEnumerator();
		}

		#endregion

		public void Add(TKey key, TValue value)
		{
			_list.Add(new KeyValuePair<TKey, TValue>(key, value));
		}
	}
}